#include <rclcpp/rclcpp.hpp>

// MoveIt Servo
#include <moveit_servo/servo_parameters.h>
#include <moveit_servo/servo.h>
#include <moveit/planning_scene_monitor/planning_scene_monitor.h>

#include <control_msgs/msg/joint_jog.hpp>
#include <geometry_msgs/msg/twist_stamped.hpp>
#include <geometry_msgs/msg/wrench_stamped.hpp>
#include <tf2_ros/buffer.h>

#include <rcl_interfaces/msg/set_parameters_result.hpp>  // rqt/CLI 动态参数
#include <atomic>
#include <algorithm>
#include <cmath>
#include <string>

using namespace std::chrono_literals;

static const rclcpp::Logger LOGGER =
    rclcpp::get_logger("smc0.servo_node.cpp");

/// ============================ 说明（要点） =================================
/// 本文件基于你的原始代码，将 publishCommands() 中的 PID 恒力控制替换为 ADRC：
/// 1) TD（Tracking Differentiator）安排设定力过渡（保留你的实现）
/// 2) ESO（扩张状态观测器）估计力、力导数与总扰动
/// 3) 误差反馈律（线性PID 或 非线性 fhan，可在线切换）
/// 4) 阻抗模型 m_d x¨ + b_d x˙ + k_d x = u  将控制量 u 转换为末端位移修正，再输出速度 x˙
/// 5) 输出限幅（加速度与速度）
///
/// 在线参数（可 ros2 param set 调节，详见 main() 中 declare_parameter 部分）：
/// f_target                目标恒力（N）
/// axis                    速度输出轴向（0:x, 1:y, 2:z）
/// max_vel                 末端线速度上限（m/s）
/// limit.max_acc           末端加速度上限（m/s^2）
/// td_enable, td_r0        TD开关与“最速”参数
/// eso.beta01/02/03, eso.delta  ESO增益与 fal 小误差阈值
/// adrc.use_nlf            误差反馈律是否采用非线性（true=fhan; false=线性 PID）
/// adrc.kp/ki/kd           线性误差反馈（PID）增益
/// adrc.beta0              非线性误差反馈中的积分权重（u = beta0*∫e - fhan(...)）
/// adrc.fhan_r             非线性误差反馈中 fhan 速率参数 r
/// imp.md/bd/kd            虚拟阻抗参数（质量/阻尼/刚度）
/// ========================================================================

/// ========== 小工具：避免 C++17 对 std::clamp 的依赖 ==========
static inline double clamp_scalar(double v, double lo, double hi) {
  return v < lo ? lo : (v > hi ? hi : v);
}
static inline double sgn(double x) {
  return (x > 0) - (x < 0);
}

/// ===========================
/// 节点/通信句柄
/// ===========================
rclcpp::Node::SharedPtr node_;
rclcpp::Publisher<geometry_msgs::msg::TwistStamped>::SharedPtr twist_cmd_pub_;
rclcpp::Subscription<geometry_msgs::msg::WrenchStamped>::SharedPtr sim_force_sub_;
rclcpp::node_interfaces::OnSetParametersCallbackHandle::SharedPtr g_param_cb_handle;

/// ===========================
/// 目标恒力 & 轴向选择
/// ===========================
std::atomic<double> g_f_target{5.0};   // 目标恒力，默认 5N
std::atomic<int>    g_axis{0};         // 速度输出轴向：0=x, 1=y, 2=z

/// ===========================
/// 速度/加速度限幅
/// ===========================
std::atomic<double> g_max_vel{0.5};    // 末端线速度上限（m/s）
std::atomic<double> g_max_acc{1.5};    // 末端加速度上限（m/s^2）

/// ===========================
/// TD（Tracking Differentiator）相关（保留你的实现）
/// 将设定力 F_target 平滑为 v1（设定力） 与 v2（其导数）
/// ===========================
std::atomic<bool>   g_td_enable{true};
std::atomic<double> g_td_r0{5.0};
double g_td_v1 = 0.0;      // TD 状态 v1（平滑后的设定力）
double g_td_v2 = 0.0;      // TD 状态 v2（设定力导数）
bool   g_td_inited = false;

// Han 型 TD 的核心“fhan”计算
static inline double fhan(double x1_minus_ref, double x2, double r, double h)
{
  const double d  = r * h * h;
  const double a0 = h * x2;
  const double y  = x1_minus_ref + a0;
  const double abs_y = std::fabs(y);
  const double a1 = std::sqrt(d * (d + 8.0 * abs_y));
  const double a2 = a0 + sgn(y) * (a1 - d) * 0.5;
  const double a  = a2;
  double f = -r * ( (std::fabs(a) > d) ? sgn(a) : (a / d) );
  return f;
}
static inline void td_update(double Ref, double r0, double h, double &v1, double &v2)
{
  const double f0 = fhan(v1 - Ref, v2, r0, h);
  v1 += h * v2;
  v2 += h * f0;
}

/// ===========================
/// 力测量
/// ===========================
std::atomic<double> g_latest_fz{0.0};  // 最近一次传感器 Z 向力（示例：来自 /sim_force）
void simForceCallback(const geometry_msgs::msg::WrenchStamped::SharedPtr msg)
{
  // 本例以 Fz 为被控量，如需改用 Fx/Fy 可在此处切换
  g_latest_fz.store(msg->wrench.force.z, std::memory_order_relaxed);
  RCLCPP_INFO_THROTTLE(LOGGER, *node_->get_clock(), 500,
                       "[/sim_force] Fz=%.3f N (target=%.2f)",
                       msg->wrench.force.z, g_f_target.load());
}

/// ===========================
/// ADRC：ESO + 误差反馈 + 阻抗模型
/// ===========================

/// --- ESO 参数 ---
std::atomic<double> g_eso_beta01{60.0};
std::atomic<double> g_eso_beta02{600.0};
std::atomic<double> g_eso_beta03{2000.0};
std::atomic<double> g_eso_delta{0.01};

/// ESO 状态：z1=观测力，z2=观测力导数，z3=总扰动估计
double g_z1 = 0.0, g_z2 = 0.0, g_z3 = 0.0;
bool   g_eso_inited = false;

/// fal 非线性函数（小误差时线性化，提高 ESO 灵敏度）
static inline double fal(double e, double alpha, double delta)
{
  const double abse = std::fabs(e);
  if (abse <= delta) return e / std::pow(delta, 1.0 - alpha);
  return sgn(e) * std::pow(abse, alpha);
}

/// --- 误差反馈律参数（线性/非线性可选） ---
std::atomic<bool>   g_use_nlf{true};  // true=非线性 fhan；false=线性（PID式）
std::atomic<double> g_err_beta0{0.0}; // 非线性时的积分权重 beta0
// 线性（PID式）增益（独立于旧PID，避免混淆）
std::atomic<double> g_adrc_kp{0.8};
std::atomic<double> g_adrc_ki{0.0};
std::atomic<double> g_adrc_kd{0.0};
// 非线性 fhan 速率参数
std::atomic<double> g_err_fhan_r{50.0};

/// 误差积分 e0（用于稳态消除）
double g_e0 = 0.0;
constexpr double kE0Max = 2.0;   // 可视需要调大

/// 上一周期控制量（ESO 更新需要 u(k)）
double g_u_prev = 0.0;

/// --- 阻抗模型参数 ---
std::atomic<double> g_md{1.0};        // 虚拟质量
std::atomic<double> g_bd{40.0};       // 虚拟阻尼
std::atomic<double> g_imp_kd{200.0};  // 虚拟刚度（注意不是“微分”项） ←← 重命名避免冲突

/// 阻抗状态：x（位移修正），x_dot（速度）；输出的就是 x_dot
double g_x = 0.0, g_x_dot = 0.0;

/// 时间戳
rclcpp::Time g_prev_time;
bool g_has_prev_time = false;

/// ===========================
/// （兼容保留）旧 PID 参数占位 —— 不用于 ADRC，仅为兼容命令行参数
/// ===========================
std::atomic<double> g_pid_kp{0.05};
std::atomic<double> g_pid_ki{0.00};
std::atomic<double> g_pid_kd{0.00};

/// ===========================
/// 定时器回调：ADRC 计算并发布 Twist（含 TD/ESO/NLEF/阻抗/限幅）
/// ===========================
void publishCommands()
{
  // 计算 dt
  const rclcpp::Time now = node_->now();
  double dt = 0.01;  // 默认 10ms
  if (g_has_prev_time) {
    dt = std::max(1e-4, (now - g_prev_time).seconds());
  }
  g_prev_time = now;
  g_has_prev_time = true;

  // 读取当前测量力（例：Fz）
  const double Fe = g_latest_fz.load(std::memory_order_relaxed);

  // --- TD：把设定力平滑为 v1（首次用当前测量初始化，避免跳变） ---
  if (!g_td_inited) {
    g_td_v1 = Fe;     // 平滑从实际力缓启动到目标
    g_td_v2 = 0.0;
    g_td_inited = true;
  }
  const double F_ref_raw = g_f_target.load(std::memory_order_relaxed);
  double F_ref_eff = F_ref_raw;
  if (g_td_enable.load(std::memory_order_relaxed)) {
    const double r0 = g_td_r0.load(std::memory_order_relaxed);
    td_update(F_ref_raw, r0, dt, g_td_v1, g_td_v2);
    F_ref_eff = g_td_v1;
  }

  // --- ESO 初始化（首周期按当前测量置初值） ---
  if (!g_eso_inited) {
    g_z1 = Fe; g_z2 = 0.0; g_z3 = 0.0;
    g_u_prev = 0.0;
    g_eso_inited = true;
  }

  // --- ESO 更新（离散形式） ---
  const double beta01 = g_eso_beta01.load();
  const double beta02 = g_eso_beta02.load();
  const double beta03 = g_eso_beta03.load();
  const double delta  = g_eso_delta.load();
  const double e_obs  = g_z1 - Fe;               // 观测误差：z1 - Fe
  const double fal_e  = fal(e_obs, 0.5, delta);  // 非线性小误差函数

  // z1, z2, z3 递推；u(k) 使用上一周期控制量 g_u_prev
  g_z1 = g_z1 + dt * ( g_z2 - beta01 * e_obs );
  g_z2 = g_z2 + dt * ( g_z3 - beta02 * fal_e + g_u_prev );
  g_z3 = g_z3 + dt * ( - beta03 * fal_e );

  // --- 力跟踪误差 ---
  const double e1 = F_ref_eff - g_z1;   // 力误差（设定-观测）
  const double e2 = g_td_v2  - g_z2;    // 力误差导数（设定导-观测导）
  g_e0 += e1 * dt;                      // 误差积分
  g_e0 = clamp_scalar(g_e0, -kE0Max, kE0Max);

  // --- 误差反馈律：线性 or 非线性（fhan） ---
  double u_control = 0.0;
  if (g_use_nlf.load()) {
    // 非线性：u = beta0*e0 - fhan(e1, e2, r, dt)
    const double beta0 = g_err_beta0.load();
    const double r1    = g_err_fhan_r.load();
    const double u_fhan = fhan(e1, e2, r1, dt);
    u_control = beta0 * g_e0 - u_fhan;
  } else {
    // 线性：u = Kp*e1 + Ki*e0 + Kd*e2
    const double Kp = g_adrc_kp.load();
    const double Ki = g_adrc_ki.load();
    const double Kd = g_adrc_kd.load();
    u_control = Kp * e1 + Ki * g_e0 + Kd * e2;
  }

  // 供下周期 ESO 使用
  g_u_prev = u_control;

  // --- 阻抗模型：m_d x¨ + b_d x˙ + k_d x = u ---
  const double md = std::max(1e-6, g_md.load());
  const double bd = g_bd.load();
  const double kd = g_imp_kd.load();  // ←← 使用重命名后的虚拟刚度

  double x_dd = (u_control - bd * g_x_dot - kd * g_x) / md;

  // --- 输出限幅：加速度与速度 ---
  const double Amax = std::max(0.0, g_max_acc.load());
  if (x_dd >  Amax) x_dd =  Amax;
  if (x_dd < -Amax) x_dd = -Amax;

  double new_x_dot = g_x_dot + x_dd * dt;
  const double Vmax = std::max(0.0, g_max_vel.load());
  if (new_x_dot >  Vmax) new_x_dot =  Vmax;
  if (new_x_dot < -Vmax) new_x_dot = -Vmax;

  g_x_dot = new_x_dot;
  g_x    += g_x_dot * dt;

  // --- 发布 TwistStamped 指令（将 x˙ 映射到所选轴向） ---
  auto msg = std::make_unique<geometry_msgs::msg::TwistStamped>();
  msg->header.stamp = now;
  msg->header.frame_id = "base_link";

  const int axis = g_axis.load();
  msg->twist.linear.x = (axis == 0) ? g_x_dot : 0.0;  // 默认为 X 轴（与你原始代码一致）
  msg->twist.linear.y = (axis == 1) ? g_x_dot : 0.0;
  msg->twist.linear.z = (axis == 2) ? g_x_dot : 0.0;
  msg->twist.angular.x = 0.0;
  msg->twist.angular.y = 0.0;
  msg->twist.angular.z = 0.0;

  twist_cmd_pub_->publish(std::move(msg));

  // --- 关键量节流打印 ---
  RCLCPP_INFO_THROTTLE(
      LOGGER, *node_->get_clock(), 50,
      "[ADRC] Fe=%.3fN  Fref=%.3fN  e1=%.3f  e2=%.3f  u=%.3f  x_dot=%.4f (dt=%.4f)"
      " [ESO z1=%.3f z2=%.3f z3=%.3f | axis=%d Vmax=%.3f Amax=%.3f | TD:%s r0=%.1f v2=%.3f]",
      Fe, F_ref_eff, e1, e2, g_u_prev, g_x_dot, dt,
      g_z1, g_z2, g_z3, axis, Vmax, Amax,
      g_td_enable.load() ? "ON" : "OFF",
      g_td_r0.load(), g_td_v2);
}

int main(int argc, char** argv)
{
  rclcpp::init(argc, argv);
  rclcpp::NodeOptions node_options;
  node_options.use_intra_process_comms(false);
  node_ = std::make_shared<rclcpp::Node>("servo_node", node_options);

  // 等待 RViz/环境就绪
  rclcpp::sleep_for(std::chrono::seconds(4));

  auto tf_buffer = std::make_shared<tf2_ros::Buffer>(node_->get_clock());
  auto planning_scene_monitor = std::make_shared<planning_scene_monitor::PlanningSceneMonitor>(
      node_, "robot_description", tf_buffer, "planning_scene_monitor");

  if (planning_scene_monitor->getPlanningScene())
  {
    planning_scene_monitor->startStateMonitor("/joint_states");
    planning_scene_monitor->setPlanningScenePublishingFrequency(25);
    planning_scene_monitor->startPublishingPlanningScene(
        planning_scene_monitor::PlanningSceneMonitor::UPDATE_SCENE,
        "/moveit_servo/publish_planning_scene");
    planning_scene_monitor->startSceneMonitor();
    planning_scene_monitor->providePlanningSceneService();
  }
  else
  {
    RCLCPP_ERROR(LOGGER, "Planning scene 未正确配置");
    return EXIT_FAILURE;
  }

  // Publisher/Subscriber
  twist_cmd_pub_ = node_->create_publisher<geometry_msgs::msg::TwistStamped>(
      "servo_node/delta_twist_cmds", 10);

  sim_force_sub_ = node_->create_subscription<geometry_msgs::msg::WrenchStamped>(
      "/sim_force",
      rclcpp::SensorDataQoS(),
      &simForceCallback);

  // MoveIt Servo
  auto servo_parameters = moveit_servo::ServoParameters::makeServoParameters(node_);
  auto servo = std::make_unique<moveit_servo::Servo>(node_, servo_parameters, planning_scene_monitor);
  servo->start();

  // ================= 参数声明（带默认值） =================
  node_->declare_parameter<double>("f_target", g_f_target.load());
  node_->declare_parameter<int>("axis", g_axis.load());
  node_->declare_parameter<double>("max_vel", g_max_vel.load());
  node_->declare_parameter<double>("limit.max_acc", g_max_acc.load());

  node_->declare_parameter<bool>("td_enable", g_td_enable.load());
  node_->declare_parameter<double>("td_r0", g_td_r0.load());

  node_->declare_parameter<double>("eso.beta01", g_eso_beta01.load());
  node_->declare_parameter<double>("eso.beta02", g_eso_beta02.load());
  node_->declare_parameter<double>("eso.beta03", g_eso_beta03.load());
  node_->declare_parameter<double>("eso.delta",  g_eso_delta.load());

  node_->declare_parameter<bool>("adrc.use_nlf", g_use_nlf.load());
  node_->declare_parameter<double>("adrc.kp",    g_adrc_kp.load());
  node_->declare_parameter<double>("adrc.ki",    g_adrc_ki.load());
  node_->declare_parameter<double>("adrc.kd",    g_adrc_kd.load());
  node_->declare_parameter<double>("adrc.beta0", g_err_beta0.load());
  node_->declare_parameter<double>("adrc.fhan_r", g_err_fhan_r.load());

  node_->declare_parameter<double>("imp.md", g_md.load());
  node_->declare_parameter<double>("imp.bd", g_bd.load());
  node_->declare_parameter<double>("imp.kd", g_imp_kd.load());  // ←← 参数键不变

  // 兼容旧 PID 参数（不再用于 ADRC，仅保留）
  node_->declare_parameter<double>("kp", g_pid_kp.load());
  node_->declare_parameter<double>("ki", g_pid_ki.load());
  node_->declare_parameter<double>("kd", g_pid_kd.load());

  // 同步初值
  g_f_target.store(node_->get_parameter("f_target").as_double());
  g_axis.store(node_->get_parameter("axis").as_int());
  g_max_vel.store(node_->get_parameter("max_vel").as_double());
  g_max_acc.store(node_->get_parameter("limit.max_acc").as_double());

  g_td_enable.store(node_->get_parameter("td_enable").as_bool());
  g_td_r0.store(node_->get_parameter("td_r0").as_double());

  g_eso_beta01.store(node_->get_parameter("eso.beta01").as_double());
  g_eso_beta02.store(node_->get_parameter("eso.beta02").as_double());
  g_eso_beta03.store(node_->get_parameter("eso.beta03").as_double());
  g_eso_delta.store(node_->get_parameter("eso.delta").as_double());

  g_use_nlf.store(node_->get_parameter("adrc.use_nlf").as_bool());
  g_adrc_kp.store(node_->get_parameter("adrc.kp").as_double());
  g_adrc_ki.store(node_->get_parameter("adrc.ki").as_double());
  g_adrc_kd.store(node_->get_parameter("adrc.kd").as_double());
  g_err_beta0.store(node_->get_parameter("adrc.beta0").as_double());
  g_err_fhan_r.store(node_->get_parameter("adrc.fhan_r").as_double());

  g_md.store(node_->get_parameter("imp.md").as_double());
  g_bd.store(node_->get_parameter("imp.bd").as_double());
  g_imp_kd.store(node_->get_parameter("imp.kd").as_double());

  // 旧 PID（不用于 ADRC）
  g_pid_kp.store(node_->get_parameter("kp").as_double());
  g_pid_ki.store(node_->get_parameter("ki").as_double());
  g_pid_kd.store(node_->get_parameter("kd").as_double());

  // 动态参数回调
  g_param_cb_handle = node_->add_on_set_parameters_callback(
    [](const std::vector<rclcpp::Parameter>& params)
    {
      rcl_interfaces::msg::SetParametersResult result;
      result.successful = true;
      for (const auto& p : params)
      {
        const auto &name = p.get_name();

        // 基本
        if (name == "f_target")             { g_f_target.store(p.as_double()); }
        else if (name == "axis")            { g_axis.store(p.as_int()); }
        else if (name == "max_vel")         { g_max_vel.store(p.as_double()); }
        else if (name == "limit.max_acc")   { g_max_acc.store(p.as_double()); }

        // TD
        else if (name == "td_enable")       { g_td_enable.store(p.as_bool()); }
        else if (name == "td_r0")           { g_td_r0.store(p.as_double()); }

        // ESO
        else if (name == "eso.beta01")      { g_eso_beta01.store(p.as_double()); }
        else if (name == "eso.beta02")      { g_eso_beta02.store(p.as_double()); }
        else if (name == "eso.beta03")      { g_eso_beta03.store(p.as_double()); }
        else if (name == "eso.delta")       { g_eso_delta.store(p.as_double()); }

        // ADRC 误差反馈
        else if (name == "adrc.use_nlf")    { g_use_nlf.store(p.as_bool()); }
        else if (name == "adrc.kp")         { g_adrc_kp.store(p.as_double()); }
        else if (name == "adrc.ki")         { g_adrc_ki.store(p.as_double()); }
        else if (name == "adrc.kd")         { g_adrc_kd.store(p.as_double()); }
        else if (name == "adrc.beta0")      { g_err_beta0.store(p.as_double()); }
        else if (name == "adrc.fhan_r")     { g_err_fhan_r.store(p.as_double()); }

        // 阻抗参数
        else if (name == "imp.md")          { g_md.store(p.as_double()); }
        else if (name == "imp.bd")          { g_bd.store(p.as_double()); }
        else if (name == "imp.kd")          { g_imp_kd.store(p.as_double()); } // ←← 更新点

        // 旧 PID（兼容保留）
        else if (name == "kp")              { g_pid_kp.store(p.as_double()); }
        else if (name == "ki")              { g_pid_ki.store(p.as_double()); }
        else if (name == "kd")              { g_pid_kd.store(p.as_double()); }
      }
      return result;
    }
  );

  // 控制循环：10ms（100Hz）
  rclcpp::TimerBase::SharedPtr timer = node_->create_wall_timer(10ms, publishCommands);

  auto executor = std::make_unique<rclcpp::executors::MultiThreadedExecutor>();
  executor->add_node(node_);
  executor->spin();

  rclcpp::shutdown();
  return 0;
}
